if '__main__' == __name__:
    import argparse
    import os
    import pickle
    import time
    import pydoc
    from tvts import tvts, DEFAULT_HOST, DEFAULT_PORT, DEFAULT_DB_NAME, DEFAULT_TABLE_PREFIX, DEFAULT_SAVE_FREQ

    MS_SLEEP_TO_WAIT = 5

    VER = 'v1.0'
    BASE_DIR, FILE_NAME = os.path.split(__file__)
    x_tmp_store_dir = os.path.join(BASE_DIR, '_save', FILE_NAME, VER)
    os.makedirs(x_tmp_store_dir, exist_ok=True)

    x_types = ('init', 'save_epoch', 'save_batch')

    parser = argparse.ArgumentParser()
    parser.add_argument('name', help='name of the training')
    parser.add_argument('pid', help='pid of the cli invoker', type=int)
    parser.add_argument('type', help=f'type of the cli invocation, one of {x_types}', type=str)
    parser.add_argument('args', help='arguments, 3 args a unit, i.e. name, type, value', metavar='arg', type=str, nargs='*')

    parser.add_argument('--host', help='host of the mongodb', type=str, default=DEFAULT_HOST)
    parser.add_argument('-p', '--port', help='port of the mongodb', type=int, default=DEFAULT_PORT)
    parser.add_argument('--db', help='db name of the mongodb', type=str, default=DEFAULT_DB_NAME)
    parser.add_argument('--prefix', help='table name prefix of the mongodb', type=str, default=DEFAULT_TABLE_PREFIX)

    args = parser.parse_args()

    x_name = args.name
    x_pid = args.pid
    x_type = args.type
    x_args_list = args.args

    x_host = args.host
    x_port = args.port
    x_db = args.db
    x_prefix = args.prefix

    if x_type not in x_types:
        raise Exception(f'{x_type} is not in {x_types}!')

    # collect arguments
    x_args_dict = {}
    x_n_args = len(x_args_list)
    x_n_units = (x_n_args - 1) // 3 + 1
    x_arg_units = [(x_args_list[i*3], x_args_list[i*3 + 1], x_args_list[i*3 + 2]) for i in range(x_n_units)]
    for name, type, val in x_arg_units:
        type = pydoc.locate(type)
        val = type(val)
        x_args_dict[name] = val

    name_prefix = f'{x_name}_{x_pid}'
    pkl_name = f'{name_prefix}.pkl'
    x_pkl_path = os.path.join(x_tmp_store_dir, pkl_name)
    save_path_name = f'{name_prefix}.txt'
    x_save_path_path = os.path.join(x_tmp_store_dir, save_path_name)

    if x_type == 'init':
        save_freq = int(x_args_dict.get('save_freq', DEFAULT_SAVE_FREQ))
        ts = tvts(x_name, x_host, x_port, x_db, x_prefix, x_args_dict, save_freq)
        x_train_id = x_args_dict.get('parent_id', None)
        if x_train_id is None or x_train_id == 0:
            path = ''
        else:
            x_epoch = x_args_dict.get('parent_epoch', None)
            if x_epoch is None:
                path = ts.resume(x_train_id)
            else:
                path = ts.resume(x_train_id, x_epoch)

        with open(x_pkl_path, 'wb') as f:
            ts.before_save_to_pickle()
            pickle.dump(ts, f)
        with open(x_save_path_path, 'w', encoding='utf8') as f:
            f.write(path)
    else:
        while not os.path.exists(x_save_path_path):
            time.sleep(0.001 * MS_SLEEP_TO_WAIT)
        with open(x_pkl_path, 'rb') as f:
            ts = pickle.load(f)
            ts.after_load_from_pickle()
        if x_type == 'save_epoch':
            x_epoch = int(x_args_dict['epoch'])
            save_path = x_args_dict.get('save_path', None)
            ts.save_epoch(x_epoch, x_args_dict, save_path)
        elif x_type == 'save_batch':
            x_epoch = int(x_args_dict['epoch'])
            x_batch = int(x_args_dict['batch'])
            x_is_batch_global = x_args_dict.get('is_batch_global', False)
            ts.save_batch(x_epoch, x_batch, x_args_dict, x_is_batch_global)
